license = '''
* Copyright (c) 2012, Yang Liu & Niansheng Zhou
* All rights reserved.
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of Yang Liu & Niansheng Zhou nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS "AS IS" AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''

import os, string, cPickle, time, math
import sys
import asynchat,asyncore
import socket
import threading

def BUS(i,j):
    return True

def SWITCH(i,j):
    return True

def MESH1(p):
    return lambda i,j,p=p: (i-j)**2==1

def TORUS1(p):
    return lambda i,j,p=p: (i-j+p)%p==1 or (j-i+p)%p==1

def MESH2(p):
    q=int(math.sqrt(p)+0.1)
    return lambda i,j,q=q: ((i%q-j%q)**2,(i/q-j/q)**2) in [(1,0),(0,1)]

def TORUS2(p):
    q=int(math.sqrt(p)+0.1)
    return lambda i,j,q=q: ((i%q-j%q+q)%q,(i/q-j/q+q)%q) in [(0,1),(1,0)] or \
                           ((j%q-i%q+q)%q,(j/q-i/q+q)%q) in [(0,1),(1,0)]
def TREE(i,j):
    return i==int((j-1)/2) or j==int((i-1)/2)

class Client(asynchat.async_chat):
    def __init__(self,host):
        asynchat.async_chat.__init__(self)
        self.host = host
        self.set_terminator(None)
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.connect((self.host.ROOT, self.host.PORT))
        
    def handle_close(self):
        self.close()
 
    def handle_expt(self):
        self.close()
            
    def collect_incoming_data(self, data):
        if data[0:4] == 'rank':
            self.host.rank = int(data[5:])
#            print 'client rank=',self.host.rank
        else: 
            comm = self.host
            sender_rank = int(data[0:3])
            receiver_rank = int(data[3:6])
            while comm.pipe_read_locks[sender_rank,receiver_rank]:
                time.sleep(0.1)
            comm.pipes[sender_rank,receiver_rank] = data
            comm.pipe_read_locks[sender_rank,receiver_rank] = True

class ServerChannel(asynchat.async_chat):
    def __init__(self,conn,server):
        asynchat.async_chat.__init__(self, conn)
        self.server = server
        self.set_terminator(None)
        
    def handle_close(self):
        self.close()
 
    def handle_expt(self):
        self.close()
            
    def collect_incoming_data(self, data):
        sender_rank = int(data[0:3])
        receiver_rank = int(data[3:6])
        if receiver_rank == 0:
            comm = self.server.host
            while comm.pipe_read_locks[sender_rank,0]:
                time.sleep(0.1)
            comm.pipes[sender_rank,0] = data
            comm.pipe_read_locks[sender_rank,0] = True
        else:
            sc = self.server.rank_channels[receiver_rank]
            sc.push(data)
        
class Server(asyncore.dispatcher):
    def __init__(self,ip,port,listen_port,host):
        asyncore.dispatcher.__init__(self)
        self.host = host
        self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
        self.set_reuse_addr()
        self.bind((ip, port))
        self.listen(listen_port)
        self.rank_channels = {}
        
    def handle_accept(self):
        conn, addr = self.accept()
        sc = ServerChannel(conn, self)
        client_rank = len(self.rank_channels) + 1
        sc.push('rank=%d'%client_rank)
        self.rank_channels[client_rank] = sc

class PSimThread(threading.Thread):
    def log(self,message):
        """
        logs the message into self._logfile
        """
        #print message
        if self.logfile!=None:
            self.logfile.write(message)
            
    def getNodes(self):
        f = file('Config.txt')
        lines = f.readlines()
        nodes = []
        for line in lines:
            line = line.replace(' ','')
            if line != '':
                nodes.append(line)   
        return nodes

    def __init__(self,p,topology=SWITCH,logfilename=None):
        """
        forks p-1 processes and creates p*p
        """
        threading.Thread.__init__(self)
        self.logfile = logfilename and open(logfilename,'w')
        self.topology = topology
        self.log("START: creating %i parallel processes\n" % p)
        self.nprocs = p
        self.pipes = {}
        self.pipe_read_locks = {}
        for i in range(p):
            for j in range(p):
                self.pipes[i,j] = None
                self.pipe_read_locks[i,j] = False
        nodes = self.getNodes()
        self.ROOT,self.PORT = nodes[0].split(':')
        self.PORT = int(self.PORT)
        self.LISTENING_PORT = 5051
        ip_list = socket.gethostbyname_ex(socket.gethostname())[2]
        if self.ROOT in ip_list and sys.argv[-1]!='client':
            self.rank = 0
            self.instance = Server(self.ROOT, self.PORT, self.LISTENING_PORT,self)
        else:
            self.rank = -1
            self.instance = Client(self)
        self.log("START: done.\n")

    def packData(self,sender_rank,receiver_rank,data):
        packed_data = string.zfill(str(sender_rank),3)
        packed_data += string.zfill(str(receiver_rank),3)
        s = cPickle.dumps(data)
        packed_data += string.zfill(str(len(s)),10)
        packed_data += s
        return packed_data
    
    def unpackData(self,data):
        sender_rank = int(data[:3])
        receiver_rank = int(data[3:6])
        data_len = int(data[6:16])
        raw_data = cPickle.loads(data[16:data_len+16])
        return raw_data

    def _send(self,j,data):
        """
        sends data to process #j
        """
        if j<0 or j>=self.nprocs:
            self.log("process %i: send(%i,...) failed!\n" % (self.rank,j))
            raise Exception
        self.log("process %i: send(%i,%s) starting...\n" % \
                 (self.rank,j,repr(data)))
        packed_data = self.packData(self.rank,j,data)
        if self.rank == 0:
            if j == 0:
                while self.pipe_read_locks[0,0]:
                    time.sleep(0.1)
                self.pipes[0,0] = packed_data
                self.pipe_read_locks[0,0] = True
            else:
                sc = self.instance.rank_channels[j]
                sc.push(packed_data)
        else:
            self.instance.push(packed_data)
        self.log("process %i: send(%i,%s) success.\n" % \
                 (self.rank,j,repr(data)))

    def send(self,j,data):
        if not self.topology(self.rank,j):
            raise RuntimeError, 'topology violation'
        self._send(j,data)

    def _recv(self,j):
        """
        returns the data recvd from process #j
        """
        if j<0 or j>=self.nprocs:
            self.log("process %i: recv(%i) failed!\n" % (self.rank,j))
            raise RuntimeError
        self.log("process %i: recv(%i) starting...\n" % (self.rank,j))
        while not self.pipe_read_locks[j,self.rank]:
            time.sleep(0.1)
        data = self.unpackData(self.pipes[j,self.rank])
        self.log("process %i: recv(%i) done.\n" % (self.rank,j))
        self.pipes[j,self.rank] = None
        self.pipe_read_locks[j,self.rank] = False
        return data

    def recv(self,j):
        if not self.topology(self.rank,j):
            raise RuntimeError, 'topology violation'
        return self._recv(j)

    def one2all_broadcast(self, source, value):
        self.log("process %i: BEGIN one2all_broadcast(%i,%s)\n" % \
                 (self.rank,source, repr(value)))
        if self.rank==source:
            for i in range(0, self.nprocs):
                if i!=source:
                    self._send(i,value)
        else:
            value=self._recv(source)
        self.log("process %i: END one2all_broadcast(%i,%s)\n" % \
                 (self.rank,source, repr(value)))
        return value

    def all2all_broadcast(self, value):
        self.log("process %i: BEGIN all2all_broadcast(%s)\n" % \
                 (self.rank, repr(value)))
        vector=self.all2one_collect(0,value)
        vector=self.one2all_broadcast(0,vector)
        self.log("process %i: END all2all_broadcast(%s)\n" % \
                 (self.rank, repr(value)))
        return vector

    def one2all_scatter(self,source,data):
        self.log('process %i: BEGIN all2one_scatter(%i,%s)\n' % \
                 (self.rank,source,repr(data)))
        if self.rank==source:
             h, reminder = divmod(len(data),self.nprocs)
             if reminder: h+=1
             for i in range(self.nprocs):
                 self._send(i,data[i*h:i*h+h])
        vector = self._recv(source)
        self.log('process %i: END all2one_scatter(%i,%s)\n' % \
                 (self.rank,source,repr(data)))
        return vector

    def all2one_collect(self,destination,data):
        self.log("process %i: BEGIN all2one_collect(%i,%s)\n" % \
                 (self.rank,destination,repr(data)))
        self._send(destination,data)
        if self.rank==destination:
            vector = [self._recv(i) for i in range(self.nprocs)]
        else:
            vector = []
        self.log("process %i: END all2one_collect(%i,%s)\n" % \
                 (self.rank,destination,repr(data)))
        return vector

    def all2one_reduce(self,destination,value,op=lambda a,b:a+b):
        self.log("process %i: BEGIN all2one_reduce(%s)\n" % \
                 (self.rank,repr(value)))
        self._send(destination,value)
        if self.rank==destination:
            result = reduce(op,[self._recv(i) for i in range(self.nprocs)])
        else:
            result = None
        self.log("process %i: END all2one_reduce(%s)\n" % \
                 (self.rank,repr(value)))
        return result

    def all2all_reduce(self,value,op=lambda a,b:a+b):
        self.log("process %i: BEGIN all2all_reduce(%s)\n" % \
                 (self.rank,repr(value)))
        result=self.all2one_reduce(0,value,op)
        result=self.one2all_broadcast(0,result)
        self.log("process %i: END all2all_reduce(%s)\n" % \
                 (self.rank,repr(value)))
        return result

        @staticmethod
        def sum(x,y): return x+y
        @staticmethod
        def mul(x,y): return x*y
        @staticmethod
        def max(x,y): return max(x,y)
        @staticmethod
        def min(x,y): return min(x,y)

    def barrier(self):
        self.log("process %i: BEGIN barrier()\n" % (self.rank))
        self.all2all_broadcast(0)
        self.log("process %i: END barrier()\n" % (self.rank))
        return

    def run(self):
        asyncore.loop()

class PSim():
    def __init__(self,p,topology=SWITCH):
        self.comm = PSimThread(p,topology)
        self.comm.start()
        self.wait(self.comm)
        
    def wait(self,comm):
        if comm.rank == 0:
            while len(comm.instance.rank_channels) < comm.nprocs - 1:
                time.sleep(0.5)
        else:
            while comm.rank < 0:
                time.sleep(0.1)
                
    def getComm(self):
        return self.comm
        
